#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
通用JSON数据提取API

这个模块提供了基于FastAPI的API，用于通用图像OCR识别和JSON数据提取。
接收image_file_path参数，直接返回提取的JSON数据，无需分步骤处理。
"""

import os
import sys
import json
import base64
from typing import Dict, Any, Optional, List
import time
import traceback

# 导入FastAPI相关依赖
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from py_files.invoice_prompt_ocr.api_invoice_file_type_prompt_ocr import image_to_data_url, extract_json_from_content

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py_files.config import get_client

# 创建路由器
router_json = APIRouter(
    prefix="/api/image2json",
    tags=["通用JSON数据提取"],
    responses={404: {"description": "未找到"}},
)

# 模型定义
class JsonExtractRequest(BaseModel):
    """JSON提取请求模型"""
    image_file_path: str = Field(..., description="图像文件路径")
    model_name: Optional[str] = Field(None, description="要使用的模型名称，不指定则使用默认模型。对于多模态任务，建议使用带vl或flash的模型")
    extract_type: Optional[str] = Field("通用json", description="提取类型，如'通用json'等")

# API响应模型
class JsonExtractResponse(BaseModel):
    """JSON提取响应模型"""
    status: str = Field(..., description="请求状态，如success或error")
    data: Dict[str, Any] = Field(..., description="提取的JSON数据")
    message: Optional[str] = Field(None, description="附加消息或错误信息")

def get_image_base64(image_path: str) -> str:
    """
    读取图像文件并转换为base64编码
    
    Args:
        image_path: 图像文件路径
        
    Returns:
        base64编码的图像字符串
    """
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return encoded_string
    except Exception as e:
        print(f"读取图像文件出错: {str(e)}")
        traceback.print_exc()
        raise


def extract_json_from_image_with_ocr(image_path: str, model_key: str = None, extract_type: str = "通用json") -> Dict[str, Any]:
    """
    分两步从图像中提取JSON数据：先OCR识别文本，再结构化为JSON
    
    Args:
        image_path: 图像文件路径
        model_key: 模型名称，默认为None时会使用配置中默认模型
        extract_type: 提取类型，默认为"通用json"
        
    Returns:
        包含JSON数据和OCR文本的字典
    """
    print(f"开始从图像提取JSON数据...") 
    try:
        # 读取图像文件
        data_url = image_to_data_url(image_path)
        # 构建OCR提示词
        ocr_system_prompt = (
            "你是一个专业的OCR文本识别助手，请从我提供的图片中提取所有可见的文字内容。"
            "只需返回识别到的纯文本内容，不要进行任何解释或格式化。"
            "保持原始的文本布局，每行信息独立成行。"
            "对于数字、金额和日期，请保持原始格式不变。"
            "不要漏掉任何小数点、百分号或货币符号。"
            "如果是表格，请尽量保持表格的行列关系。"
            "务必保留所有特殊符号。"
        )
        
        # 第一步：调用模型进行OCR识别
        client, model_name = get_client(model_key)
        # 准备带图像的消息
        ocr_request_messages = [
            {"role": "system", "content": ocr_system_prompt},
            {"role": "user", "content": [
                {"type": "text", "text": "请识别此图片中的所有文字内容："},
                {"type": "image_url", "image_url": data_url}
            ]}
        ]
        
        # 发送OCR请求
        ocr_response = client.chat.completions.create(
            model=model_name,
            messages=ocr_request_messages,
            max_tokens=3000,
            temperature=0.2
        )
        
        # 获取OCR文本结果
        ocr_text = ocr_response.choices[0].message.content
        print(f"OCR识别结果 (部分显示): {ocr_text[:200]}...")
        
        # 第二步：调用模型进行JSON提取
        # 构建JSON提取提示词
        json_system_prompt = (
            "你是一个专业的结构化数据提取助手。请将提供的OCR文本中的关键信息提取为规范的JSON格式，要求如下："
            "1. 输出格式必须严格遵循JSON格式，不要有任何其他文字说明。"
            "2. 根据内容自动识别字段名(key)和对应值(value)，使用最合适的字段名来表示内容。"
            "3. 根据内容类型进行不同处理："
            "   - 文档/表单：识别各个字段的标签和值，结构化提取"
            "   - 表格：使用表头作为key，单元格内容作为value"
            "   - 卡片/证件：提取各类标识信息"
            "   - 海报/广告：提取关键产品信息、价格、日期等"
            "4. 所有数值保留原始格式和单位"
            "5. 对于不存在的信息，相应字段赋值为null"
            "6. 对于特殊符号和字符，请完整保留"
        )
        
        # 准备JSON提取请求
        json_request_messages = [
            {"role": "system", "content": json_system_prompt},
            {"role": "user", "content": f"""请将以下OCR文本内容提取为JSON格式：

OCR文本内容:
{ocr_text}

请确保：
1. 仅返回JSON格式，不要有任何解释说明
2. 自动识别所有关键信息并提取为合适的字段
3. 保留所有数字格式、单位和特殊符号
4. 根据文本格式和内容智能构建JSON结构
5. 对于表格数据，尝试保持其行列关系"""
            }
        ]
        
        # 发送JSON提取请求
        json_response = client.chat.completions.create(
            model=model_name,
            messages=json_request_messages,
            max_tokens=3000,
            temperature=0.2
        )
        
        # 获取原始响应内容
        content = json_response.choices[0].message.content
        print(f"JSON提取结果 (部分显示): {content[:200]}...")
        
        # 使用extract_json_from_content函数安全解析JSON
        result = extract_json_from_content(content)
        # 添加OCR文本到结果
        result["ocr_text"] = ocr_text
        
        return result
    except Exception as e:
        print(f"从图像提取JSON数据出错: {str(e)}")
        traceback.print_exc()
        raise

@router_json.post("/extract", response_model=JsonExtractResponse)
async def extract_json_from_image(request: JsonExtractRequest) -> Dict[str, Any]:
    """
    从图像文件中提取JSON数据
    
    Args:
        request: 请求对象，包含image_file_path和可选的model_name
    
    Returns:
        提取的JSON数据，包含status和data字段
    """
    try:
        # 检查文件是否存在
        image_path = request.image_file_path
        # 替换路径中可能的斜杠问题
        image_path = image_path.replace('\\', '/')
        
        if not os.path.exists(image_path):
            # 尝试查找相对路径
            base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            alternative_path = os.path.join(base_dir, image_path)
            if os.path.exists(alternative_path):
                image_path = alternative_path
            else:
                return {
                    "status": "error",
                    "message": f"文件不存在: {image_path}，请检查路径是否正确",
                    "data": {}
                }
        
        # 获取请求中的模型名称和提取类型
        model_name = request.model_name if request.model_name and request.model_name.strip() else None
        extract_type = request.extract_type
        
        # 使用分步式方法从图像提取JSON数据
        print(f"开始从图像 {image_path} 提取JSON数据")
        result = extract_json_from_image_with_ocr(image_path, model_name, extract_type)
        
        # 构建响应数据
        response_data = {
            "raw_content": result.get("raw_content", ""),
            "ocr_text": result.get("ocr_text", "")
        }
        
        # 如果成功解析了JSON，添加到响应中
        if "parsed_json" in result:
            response_data["parsed_json"] = result["parsed_json"]
            
            # 返回成功状态的响应
            message = "成功提取图像文本并解析为JSON格式"
            if "parse_error" in result:
                message += f"，但注意可能存在解析问题: {result['parse_error']}"
                
            return {
                "status": "success",
                "data": response_data,
                "message": message
            }
        else:
            # 返回包含原始内容的响应
            return {
                "status": "success",
                "data": response_data,
                "message": "成功提取图像文本，但JSON解析失败，返回原始内容"
            }
    
    except json.JSONDecodeError as e:
        error_message = f"解析JSON失败: {str(e)}"
        print(error_message)
        
        return {
            "status": "error",
            "message": error_message,
            "data": {}
        }
    
    except Exception as e:
        error_traceback = traceback.format_exc()
        error_message = f"处理请求时出错: {str(e)}"
        print(error_message)
        print(error_traceback)
        
        return {
            "status": "error",
            "message": error_message,
            "data": {}
        }